% 浅水方程
% 浅水方程SWE是
% \pdv{U}{t} +\pdv{F}{x} +\pdv{G}{y} =0;
% 其中（这使SWE相当于三个方程的组合）：
% U=[h,hu,hv]
% F=[hu,hu^2+1/2gh^2,huv]
% G=[hu,hu^2+1/2gh^2,huv]
% h是水深(h>0), u是水的x方向速度，v是y方向速度，g是重力加速度
% SWE需要使用Lax-Wendroff法迭代
% 参考
% SYRACUSE UNIVERSITY, Shallow Water Equations, PHY 307
% FSU, NUMERICAL SOLUTION OF THE SHALLOW WATER EQUATIONS
% WHOI, the Shallow-Water Equation
% 警告！可能有bug！请勿用于正式用途！
% 使用AI辅助
% Gitee Repo

clc
clear

function [F,G] = compute_FG(h,u,v)
    global g

    F(:,:,1) = h.*u;
    F(:,:,2) = h.*u.^2+0.5*g*h.^2;
    F(:,:,3) = h.*u.*v;

    G(:,:,1) = h.*v;
    G(:,:,2) = h.*u.*v;
    G(:,:,3) = h.*v.^2+0.5*g*h.^2;

    return
end

function [h,u,v] = compute_huv(U)
    h = squeeze(U(:,:,1));
    u = squeeze(U(:,:,2)./h);
    v = squeeze(U(:,:,3)./h);
    return
end

% 反射边界
function U = applyBC(_U)
    global n

    U = _U;
    U(1,:,1) = U(2,:,1);
    U(:,1,1) = U(:,2,1);
    U(n,:,1) = U(n-1,:,1);
    U(:,n,1) = U(:,n-1,1);

    % 速度取反
    U(1,:,2) = -U(2,:,2);
    U(n,:,2) = -U(n-1,:,2);

    U(:,1,3) = -U(:,2,3);
    U(:,n,3) = -U(:,n-1,3);
    return;
end


global n
global g
global x
global y
global L

L = 25;
dx = 1;
dt = 0.01;
g = 10;

[y,x] = meshgrid(-L:dx:L,-L:dx:L);
n = size(x,1);

function h = create_droplet(x0,y0,h0)
    global x
    global y
    global L
    h = h0*exp(-0.01*((x-x0).^2+(y-y0).^2));
end

h = ones(n,n); %water height
##h += create_droplet(0,0,1);
h += create_droplet((2*rand()-1)*0.8*L,(2*rand()-1)*0.8*L,0.1*(2*rand()-1)+0.8);
h += create_droplet((2*rand()-1)*0.8*L,(2*rand()-1)*0.8*L,0.1*(2*rand()-1)+0.8);

u = zeros(n,n); % x speed
v = zeros(n,n); % y speed

U0 = zeros(n,n,3);
U0(:,:,1) = h;
U0(:,:,2) = h.*u;
U0(:,:,3) = h.*v; %U: [h,hu,hv]
U1 = zeros(n,n,3); %[h,hu,hv]
[F,G] = compute_FG(h,u,v); %F:[hu,hu^2+1/2gh^2,huv],G:[hv,huv,hv^2+1/2gh^2]

Um = zeros(n-1,n-1,3);

t = 0;
t_final = 10;
k = 0;

tick = 0;
TICK = 5000;

figure
scatter(0,0)
print('test.png')

for tick = 1:TICK
    i = 1:n-1;
    j = 1:n-1;

    % 计算中间步骤变量.
    % 若U0,U1,F,G定义在空间(i,j)以及时刻tick上，则Um,Fm,Gm定义在空间(i+1/2,j+1/2)和时刻tick+1/2上
    % 这种使用定义在半步长/时刻的变量是Lax-Wendroff方法的特点之一
    Um(i,j,:) = 1/4*(U0(i,j,:)+U0(i+1,j,:)+U0(i,j+1,:)+U0(i+1,j+1,:)) ...
         - 1/2*dt/dx * ((F(i+1,j,:)-F(i,j,:)) + (F(i+1,j+1,:)-F(i,j+1,:)))/2 ...
         - 1/2*dt/dx * ((G(i,j+1,:)-G(i,j,:)) + (G(i+1,j+1,:)-G(i+1,j,:)))/2;

    [hm,um,vm] = compute_huv(Um);
    [Fm,Gm] = compute_FG(hm,um,vm);

    i = 2:n-1;
    j = 2:n-1;

    % 根据Fm,Gm计算U1
    U1(i,j,:) = U0(i,j,:) ...
         - dt/dx * ((Fm(i,j,:)-Fm(i-1,j,:)) + (Fm(i,j-1,:)-Fm(i-1,j-1,:)))/2 ...
         - dt/dx * ((Gm(i,j,:)-Gm(i,j-1,:)) + (Gm(i-1,j,:)-Gm(i-1,j-1,:)))/2;

    U1 = applyBC(U1);
    [h,u,v] = compute_huv(U1);
    [F,G] = compute_FG(h,u,v);

    U0 = U1;

    if any(isnan(U1(:))) || any(isinf(U1(:)))
        error('diverged.')
    end

    if mod(tick,20) == 1
        clf
        hold on
        axis equal
        axis([-L,L,-L,L,0,2])
        caxis([0.5,1.5])
        xlabel('x')
        ylabel('y')
        surf(x,y,h,'Edgecolor','none')
        view(30,45)
##        view(0,90)
        drawnow
        pause(0.01)
##        if mod(tick,100) == 1
##            print(sprintf('my_%d.png',tick))
##        end
    end
end


